package com.mou.aiagent.agent.model;

import cn.hutool.core.util.StrUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

/**
 * 代理基类,步骤的定义和执行流程
 */
@Slf4j
@Data
public abstract class BaseAgent {
    // 执行状态 默认为空闲状态
    private AgentState state= AgentState.IDLE;

    // 核心属性
    private String name;

    // 提示词
    private String SystemPrompt;
    private String nextStepPrompt;

    // 1执行步骤控制
    // 1.1最小执行步骤控制
    private int CurrentStep=0;
    // 1.2最大执行步骤 10
    private int MaxStep=10;
    // LLM 模型
    private ChatClient chatClient;

    //自定义Memory记忆（自主维护上下文）
    private List<Message> messageList = new ArrayList<>();

    /**
     * 运行代理方法
     * @param usrPrompt 用户的提示词
     * @return 执行的结果
     */
    public String run (String usrPrompt){
        // 1.1检查状态
        if (this.state != AgentState.IDLE){
            throw  new  RuntimeException("Agent is not idle, current state is " +this.state);
        }
        if (StrUtil.isBlank(usrPrompt)){
            throw  new  RuntimeException("Prompt is empty");
        }
        // 1.2设置状态为运行中
        this.state = AgentState.RUNNING;
        // 1.3存储用户提示词到 messageList 中
        messageList.add(new UserMessage(usrPrompt));
        // 1.4保存结果列表
        List<String> results = new ArrayList<>();
        try {
            // 1.5循环执行步骤
            for (int i =0; i<MaxStep && state != AgentState.FINISHED; i++){
                // 1.5.1设置当前步骤
                int stepNumber = i + 1;
                // 存储
                CurrentStep = stepNumber;
                log.info("Current step: {}", stepNumber);
                // 1.5.1执行步骤
                String step = step();
                String result = stepNumber + ":" + step;
                results.add(result);
            }
            // 检查是否执行到最大步骤
            if (CurrentStep == MaxStep){
                // 1.6设置状态为已完成
                state = AgentState.FINISHED;
    //            log.info("Max step reached, agent will be terminated.");
                // 添加信息
                results.add("Max step reached, agent will be terminated.");
            }

            // 将results数组中的元素用换行符连接起来，并返回连接后的字符串
            return String.join("\n", results);
        } catch (Exception e) {
            // 1.6设置状态为错误
           state = AgentState.ERROR;
          log.info("Agent encountered an error: {}", e.getMessage());
        } finally {
            // 1.7清理
            cleanup();
        }
        return usrPrompt;
    }
    /**
     * 流式方法代理
     * @param usrPrompt 用户的提示词
     * @return 执行的结果
     */
    public SseEmitter runStream (String usrPrompt) {
        // 创建一个SseEmitter对象
        SseEmitter emitter = new SseEmitter(180000L); // 设置超时时间为3分钟
        // 异步处理,避免阻塞
        CompletableFuture.runAsync(() -> {
            try {
                // 1.1检查状态
                if (this.state != AgentState.IDLE) {
                    emitter.send("错误：无法运行代理" + this.state);
                    emitter.complete();
//                throw new RuntimeException("Agent is not idle, current state is " + this.state);
                }
                if (StrUtil.isBlank(usrPrompt)) {
                    emitter.send("错误：不能使用空提示词运行代理" + this.state);
                    emitter.complete();
                    throw new RuntimeException("Prompt is empty");
                }
            } catch (Exception e) {
                emitter.completeWithError(e);
            }

            // 1.2设置状态为运行中
            this.state = AgentState.RUNNING;
            // 1.3存储用户提示词到 messageList 中
            messageList.add(new UserMessage(usrPrompt));
            // 1.4保存结果列表
            List<String> results = new ArrayList<>();
            try {
                // 1.5循环执行步骤
                for (int i = 0; i < MaxStep && state != AgentState.FINISHED; i++) {
                    // 1.5.1设置当前步骤
                    int stepNumber = i + 1;
                    // 存储
                    CurrentStep = stepNumber;
                    log.info("Current step: {}", stepNumber);
                    // 1.5.1执行步骤
                    String step = step();
                    String result = stepNumber + ":" + step;
                    results.add(result);
                    // 1.6 发送结果 得到每一步给 SSE 到前端
                    emitter.send(result);
                }
                // 检查是否执行到最大步骤
                if (CurrentStep == MaxStep) {
                    // 1.6设置状态为已完成
                    state = AgentState.FINISHED;
                    //            log.info("Max step reached, agent will be terminated.");
                    // 添加信息
                    results.add("Max step reached, agent will be terminated.");
                    // 发送结果
                    emitter.send("执行结束，达到最大的步骤." + getMaxStep());
                }
                // 正常完成
                emitter.complete();
//                // 将results数组中的元素用换行符连接起来，并返回连接后的字符串
//                return String.join("\n", results);
            } catch (Exception e) {
                state = AgentState.ERROR;
                log.info("Agent encountered an error: {}", e.getMessage());
                try {
                    emitter.send("错误：代理遇到错误" + e.getMessage());
                } catch (IOException ex) {
                    emitter.completeWithError(ex);
                }
                emitter.complete();
            } finally {
                // 1.7清理
                cleanup();
            }

        });
        // 设置超时时间回调
        emitter.onTimeout(() -> {
            state = AgentState.ERROR;
            this.cleanup();
            log.info("Agent execution timed out.");
        });
        // 设置完成回调
        emitter.onCompletion(() -> {
            if (this.state==AgentState.FINISHED){
                state = AgentState.FINISHED;
                this.cleanup();
            }
            this.cleanup();
            log.info("Agent execution completed.");
        });
        return emitter;
    }


    /**
     * 定义单个步骤（循环）
     * @return
     */
   public abstract String step();

   protected void cleanup (){

   }
}
